""" Preprocess dataset for Code Generation such as LiveCodeBench task """

import os
import json
import argparse
import pandas as pd

from prompt import *
from tqdm import tqdm
from datasets import Dataset, load_dataset
from verl.utils.hdfs_io import copy, makedirs

# Load custom JSONL dataset
def gen_from_jsonl(path):
    with open(path) as f:
        for line in f:
            yield json.loads(line)

# one_shot = """
# Okay, the user wants me to write a Python code that outputs "Hello World." Let me think about how to approach this.

# First, I remember that in Python, the most basic way to print something is using the print() function. So I should use that.

# The string to print is "Hello World." So inside the print function, I need to put that string in quotes. Either single or double quotes should work, right? So maybe print("Hello World.") would do it.

# Wait, should there be a newline after? The print function in Python adds a newline by default, so that's fine. The user probably just wants the output without any extra formatting.

# Let me double-check. If I write print("Hello World."), when the code runs, it should output exactly that. No semicolons needed in Python, unlike some other languages.

# Is there any other way to do this? Maybe using variables, but the user didn't ask for anything complicated. The simplest solution is best here.

# So the code should be straightforward. Just one line: print("Hello World."). That's it.
# """.strip()

# Make prefix for the question
def make_prefix(example):
    # SYSTEM_PROMPT = (
    #     "A conversation between User and Assistant. The User asks a question, and the Assistant solves it. "
    #     "The Assistant first thinks about the reasoning process in the mind and then provides the user with the answer. "
    #     "The reasoning process and answer are enclosed within <think> </think> and <answer> </answer> tags, "
    #     "respectively, i.e., <think> reasoning process here </think> <answer> answer here </answer>. "
    #     "And the final code in the answer needs to be wrapped in ```, such as ```code block here```. "
    #     "The Assistant can think about anything in the reasoning process as long as it contributes to the answer. "
    #     "User: {prompt} Assistant: <think>"
    # )
    
    # SYSTEM_PROMPT = (
    #     "A conversation between User and Assistant. The User asks a question, and the Assistant solves it. "
    #     "The Assistant first thinks about the reasoning process in the mind and then provides the user with the answer. "
    #     "The Assistant needs to be constantly in the process of reflection, or self-denial, "
    #     "e.g. using 'rethink', 'try again', etc. to correct their thinking. "
    #     "The reasoning process and answer are enclosed within <think> </think> and <answer> </answer> tags, respectively, i.e., "
    #     "<think> Let's first think ... Oh we need to verify ... "
    #     "But wait, there may be some problems ... Let's try again ... "
    #     "Let me rethink about it ... Or we can recheck here ... </think> "
    #     "<answer> final answer here </answer>. User: {prompt}\nAssistant: <think> Let's first think"
    # )

    # SYSTEM_PROMPT = (
    #     "User:\nWrite a python code to output Hello World.\n"
    #     "Assistant: <think>\n{one_shot}\n</think>\n"
    #     "<answer>\n```python\nprint('Hello World.')\n``` </answer>\n"
    #     "User: {prompt}\nAssistant: <think>"
    # )

    # SYSTEM_PROMPT = ("{prompt}<think>")

    content = """[User] {shot_1_instruction}
[System] You need to show reflective behavior in your thinking such as using words like 'wait'.
[Assistant] OK! I need to reflect mutil times using words like 'wait'. <think>
{shot_1_response}

[User] {shot_2_instruction}
[System] You need to show reflective behavior in your thinking such as using words like 'wait'.
[Assistant] OK! I need to reflect mutil times using words like 'wait'. <think>
{shot_2_response}

[User] {shot_3_instruction}
[System] You need to show reflective behavior in your thinking such as using words like 'wait'.
[Assistant] OK! I need to reflect mutil times using words like 'wait'. <think>
{shot_3_response}

[User] {prompt}
[System] You need to show reflective behavior in your thinking such as using words like 'wait'.
[Assistant] OK! I need to reflect mutil times using words like 'wait'. <think>
""".format(
        shot_1_instruction=shot_1_instruction, shot_1_response=shot_1_response,
        shot_2_instruction=shot_2_instruction, shot_2_response=shot_2_response,
        shot_3_instruction=shot_3_instruction, shot_3_response=shot_3_response,
        prompt=example['problem']
    )

    # return SYSTEM_PROMPT.format(prompt=example['problem'])
    return content

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--data_path', default='/shd/fanggx/code_O1/DPO-R1/step_1_gen/35k.jsonl')
    parser.add_argument('--save_dir', default='/shd/fanggx/code_O1/verl/datasets/fewShot')
    parser.add_argument('--data_source', default='random_35k')
    parser.add_argument('--train_prop', type=float, default=0.99)
    parser.add_argument('--hdfs_dir', default=None)
    args = parser.parse_args()

    assert args.train_prop > 0 and args.train_prop < 1, '`train_prop` should be in (0, 1)'

    df = pd.DataFrame(map(json.loads, open(args.data_path, 'r').readlines()))
    df['solution'] = df['solution'].apply(json.dumps)
    raw_dataset = Dataset.from_pandas(df)
    print('>>> Raw Data Num:', len(raw_dataset))

    train_dataset = raw_dataset.select(range(int(len(raw_dataset) * args.train_prop)))
    test_dataset = raw_dataset.select(range(int(len(raw_dataset) * args.train_prop), len(raw_dataset)))

    def make_map_fn(split):
        def process_fn(example, idx):
            question = make_prefix(example)
            solution = {
                "fn_name": json.loads(example['solution']).get('fn_name', None),
                "inputs": json.loads(example['solution'])['inputs'],
                "outputs": json.loads(example['solution'])['outputs']
            }
            data = {
                "data_source": args.data_source,
                "prompt": [{
                    "role": "user",
                    "content": question,
                }],
                "ability": "code_generation",
                "reward_model": {
                    "style": "rule",
                    "ground_truth": json.dumps(solution)
                },
                "extra_info": {
                    'split': split,
                    'index': idx,
                }
            }
            return data
        return process_fn
            
    train_dataset = train_dataset.map(function=make_map_fn('train'), with_indices=True)
    test_dataset = test_dataset.map(function=make_map_fn('test'), with_indices=True)

    save_dir = args.save_dir
    hdfs_dir = args.hdfs_dir

    # Create local directory if not exists
    os.makedirs(os.path.expanduser(save_dir), exist_ok=True)

    train_dataset.to_parquet(os.path.join(save_dir, 'train.parquet'))
    test_dataset.to_parquet(os.path.join(save_dir, 'test.parquet'))

    if hdfs_dir is not None:
        makedirs(hdfs_dir)
        copy(src=local_dir, dst=hdfs_dir)